#include <sm2_sign.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

#include "sm2_ec.h"
#include "sm2_format.h"
#include "sm2_operation.h"
#include "sm2_rand.h"
#include "sm3.h"
#include "sm2_enc.h"


void print_n_bytes(uint8_t *s, size_t len)
{
    int i = 0;
    for(i = 0; i < len; ++i) {
        printf("%.2X", s[i]);
    }
    putchar('\n');
}

/**
 * @brief 瀵嗛挜娲剧敓鍑芥暟
 * 
 * @param[in] uint8_t *Z: 姣旂壒涓瞆锛岀┖闂寸敱caller鍒嗛厤
 * @param[in] const uint32_t z_len: 姣旂壒涓瞆鐨勬瘮鐗归暱搴�
 * @param[in] const uint32_t klen: 杈撳嚭瀵嗛挜姣旂壒涓茬殑姣旂壒闀垮害
 * 
 * @return KDFKeyData *
 */
KDFKeyData *KDF(const uint8_t *Z, const uint32_t z_len, const uint32_t klen)
{
    const uint32_t hash_unit_len = 32;              // Hash绠楁硶杈撳嚭闀垮害锛堝瓧鑺傦級
    uint32_t unit_len = z_len + 4;                  // #z + #ct
    uint32_t ct = 1;                                // 32bit鏋勬垚鐨勮鏁板櫒
    uint32_t floor, ceil;                           // floor, ceil鐢ㄤ簬琛ㄧず涓婁笅鐣屻��
    uint8_t flag;                                   // flag: 1->hash闀垮害涓嶈兘鏁撮櫎klen; 0->hash闀垮害鍙互鏁撮櫎klen
    uint8_t digest[32];                             // 涓存椂瀹瑰櫒锛岀敤浜庤Hash鍑芥暟浜х敓鐨勬暎鍒楀瓧鑺備覆
    uint8_t *unit = (uint8_t *)malloc(unit_len);    // 杈撳嚭K鐨勬瘡涓崟浣嶇殑涓存椂瀹瑰櫒锛岀敱浜巣_len鏃犳硶纭畾闀垮害锛屾晠鍔ㄦ�佸垎閰嶇┖闂�
    KDFKeyData *K;                                  // 杩斿洖鐨勭粨鏋滀负KDFKeyData缁撴瀯浣撴寚閽堬紝鎸囧悜杩斿洖鐨勫�兼墍鍦ㄧ┖闂�

    int t = 0, j = 0;

    // 鑾峰彇ceil(klen/v)鍜屽搴斾笅鐣�
    floor = klen / HASH_LEN;
    if(klen % HASH_LEN) {
        ceil = floor + 1;
        flag = 0;
    } else {
        ceil = floor;
        flag = 1;
    }

    t = klen - hash_unit_len * floor;               // 涓存椂鍙橀噺鐢ㄤ簬瀛樺偍鏈�鍚庝竴涓崟鍏冨簲璇ュ瓨鏀剧殑姣旂壒澶у皬锛堝崟浣嶏細bit锛�

    K = (KDFKeyData *)malloc(sizeof(KDFKeyData));
    K->k = (uint8_t *)malloc(hash_unit_len * ceil);

    memset(K->k, 0, hash_unit_len * ceil);

    memcpy(unit, Z, z_len);                         // 寰呰绠楃殑Hash鍊肩殑鍓嶉潰z_len瀛楄妭濮嬬粓涓篫

    for(ct = 1; ct < ceil; ++ct)
    {
        memset(unit+z_len, 0, 4);                   // 灏哯_ct(unit)鍚庨潰4涓瓧鑺傝缃负0
        for(j = 0; j < 4; ++j) unit[z_len+j] |= ((0xff000000 >> (j * 8)) & ct) >> ((3-j) * 8);     // Z||ct
        sm3_hash(digest, unit, unit_len);           // 鐢熸垚瑕佸～鍏ョ殑Hash鍊硷紝澶у皬涓�32瀛楄妭--256bit
        memcpy(K->k+32*(ct-1), digest, 32);         // 濉叆
    }

    memset(unit+z_len, 0, 4);                       // 灏哯_ct(unit)鍚庨潰4涓瓧鑺傝缃负0
    for(j = 0; j < 4; ++j) unit[z_len+j] |= ((0xff000000 >> (j * 8)) & ct) >> ((3-j) * 8);         // Z||ct
    sm3_hash(digest, unit, unit_len);

    printf("==========\n");
    print_n_bytes(unit, unit_len);
    print_n_bytes(digest, 32);


    // 濉叆鏈�鍚庝竴涓猽nit
    if(flag) {
        K->length = ceil*hash_unit_len*8;
        memcpy(K->k + 32*(ceil-1), digest, 32);
    }
    else {
        K->length = (ceil-1)*hash_unit_len*8+t;
        memcpy(K->k + 32*(ceil-1), digest, (uint32_t)(t/8));
        K->k[32*(ceil-1)+t/8] |= (digest[t/8] >> (t%8)) << (t%8);       // 鍏佽澶勭悊klen涓嶆槸8鐨勬暣鏁扮殑鎯呭喌
    }

    #ifdef DEBUG
    printf("KDF:\n====\n");
    printf("ceil = %d\nklen = %d\nt = %d\nk->length = %d\n", ceil, klen, t, K->length);
    printf("KDF output = ");
    print_n_bytes(K->k, K->length/8);
    printf("=======\n");
    #endif

    free(unit);
    return K;
}

/**
 * @brief SM2鍔犲瘑
 * 
 * @param[in] const char *message: 寰呭姞瀵嗙殑瀛楄妭涓�
 * @param[in] const SM2_Point *pubkey: SM2鍏挜锛屼负涓�涓豢灏勫潗鏍�
 * @param[in] const size_t message_len: 寰呭姞瀵嗗瓧鑺傜殑闀垮害
 * 
 * @return 杈撳嚭鍔犲瘑瀛楄妭涓叉寚閽�
 */
uint8_t *sm2_encrypt(uint32_t *retLength, const uint8_t *message, const uint8_t pub_key[64], const size_t message_len)
{
    sm2_uint256_t k;                                                // 闅忔満鏁�
    SM2_Jacobian_Point pub_jaco;                                    
    SM2_Jacobian_Point t_jaco;
    SM2_Point t_affine;
    SM2_Point pub;

    uint8_t *concat = malloc(64+message_len);                       // 涓棿鍙橀噺锛屽弬涓庤绠�
    memset(concat, 0, 64+message_len);

    uint8_t digest[32];

    KDFKeyData *t;
    int i = 0;
    int flag = 0;
    uint8_t *encrypted_message = NULL;                              // 鍔犲瘑鍚庣殑鏁版嵁
    encrypted_message = (uint8_t *)malloc(64+message_len+32);

    sm2_bytes_to_integer(pub.x, pub_key);
    sm2_bytes_to_integer(pub.y, pub_key+32);
    sm2_set_jacobian_point_from_std(&pub_jaco, &pub);

regen:
    #ifdef SET_RANDOM
    char_to_hex(k, "59276E27D506861A16680F3AD9C02DCCEF3CC1FA3CDBE4CE6D54B80DEAC1BC21");     // 闅忔満鏁帮紝鍥芥爣
    #endif

    // C1
    // -----------------------------------------------------------------------------------------------------------
    sm2_jacobian_point_k_mul_5naf(&t_jaco, &SM2_G, k);           // [k]G = (x1, y1)
    sm2_get_std_point(&t_affine, &t_jaco);                  // t_affine <- [k]G
    sm2_integer_to_bytes(concat, t_affine.x);
    sm2_integer_to_bytes(concat+32, t_affine.y);            // (x1||y1||<unused>)
    memcpy(encrypted_message, concat, 64);
    // -----------------------------------------------------------------------------------------------------------
    
    // S = [h]P_B锛屼絾鏄痟==1
    if(sm2_compare(pub_jaco.z, ZERO) == SM2_COMPARE_EQ)
    {
        fprintf(stderr, "鍔犲瘑閿欒: Error: S涓嶅簲璇ヤ负鏃犵┓鐐筡n");
        exit(1);
    }

    // C2
    // -----------------------------------------------------------------------------------------------------------
    sm2_jacobian_point_k_mul(&t_jaco, &pub_jaco, k);             // [k]P_{B} = (x2, y2)
    sm2_get_std_point(&t_affine, &t_jaco);                       // 寰楀埌(x2, y2)鐨勪豢灏勫潗鏍�

    // cord_concat = x2||y2
    sm2_integer_to_bytes(concat, t_affine.x);
    sm2_integer_to_bytes(concat+32, t_affine.y);                 // (x2||y2||<unused>)
    t = KDF(concat, 64, message_len*8);
    
    // 妫�鏌鏄惁涓哄叏0
    i = (t->length/8)-1;
    while(i >= 0) {
        if(t->k[i]) {
        	flag = 1;
        	break;
        }
        --i;
    }
    if(!flag) goto regen;                                       // 鑻ユ娴嬪埌t涓哄叏0鍒欓渶瑕侀噸鏂扮敓鎴愰殢鏈烘暟

    memcpy(digest, concat+32, 32);
    memcpy(concat+32+message_len, digest, 32);

    for(i = 0; i < message_len; ++i) concat[32+i] = (uint8_t)message[i] ^ t->k[i];         // (x2||M^t||y2)
    memcpy(encrypted_message+64+32, concat+32, message_len);
    // -----------------------------------------------------------------------------------------------------------

    // C3
    // -----------------------------------------------------------------------------------------------------------
    memcpy(concat+32, message, message_len);                    // (x2||message||y2)
    sm3_hash(digest, concat, 64+message_len);
    memcpy(encrypted_message+64, digest, 32);                   // 鍐欏叆C3
    // -----------------------------------------------------------------------------------------------------------
    // (C1||C3||C2)

    #ifdef DEBUG
    printf("encrypted message\n");
    print_n_bytes(encrypted_message, 64+32+message_len);
    #endif
    *retLength = 96+message_len;

    free(concat);

    return encrypted_message;
}

/**
 * @brief SM2瑙ｅ瘑
 * 
 * @param[out] uint8_t *decrypted: 瑙ｅ瘑瀛楄妭涓�
 * @param[out] uint32_t *outputLength: 瑙ｅ瘑瀛楄妭涓茬殑瀛楄妭闀垮害
 * @param[in] const uint8_t *encrypted: 寰呰В瀵嗗瓧鑺備覆(C1||C3||C2)
 * @param[in] const sm2_uint256_t *private_key: 绉侀挜锛屼负涓�涓�256bit鏁存暟
 * @param[in] const uint32_t encrypted_len: 鍔犲瘑瀛楄妭涓茬殑闀垮害
 * 
 * @return uint8_t *: 瑙ｅ瘑鍚庣殑鏄庢枃瀛楄妭涓�
 */
int sm2_decrypt(uint8_t *decrypted, uint32_t *outputLength, const uint8_t *encrypted, const uint8_t priv_key[32], const uint32_t encrypted_len)
{
	*outputLength = 0;
    int i = 0, flag = 0;
    SM2_Point t_affine;
    SM2_Jacobian_Point t_jaco;
    KDFKeyData *t = NULL;
    uint8_t cord_concat[64];
    uint8_t digest[32];
    sm2_uint256_t private_key = {0};

    uint32_t decrypted_len = encrypted_len-64-32;
    memset(decrypted, 0, decrypted_len+1);

    uint8_t *u = malloc(64+decrypted_len);

    sm2_bytes_to_integer(private_key, priv_key);

    //C1
    // -----------------------------------------------------------------------------------------------------------
    sm2_bytes_to_integer(t_affine.x, encrypted);
    sm2_bytes_to_integer(t_affine.y, encrypted+32);
    sm2_set_jacobian_point_from_std(&t_jaco, &t_affine);
    if(!sm2_point_on_curve(&t_jaco)) {
        return -1;
    }

    if(sm2_compare(t_jaco.z, ZERO) == SM2_COMPARE_EQ) {
        return -1;
    }

    sm2_jacobian_point_k_mul(&t_jaco, &t_jaco, private_key);
    sm2_get_std_point(&t_affine, &t_jaco);
    sm2_integer_to_bytes(cord_concat, t_affine.x);
    sm2_integer_to_bytes(cord_concat+32, t_affine.y);
    t = KDF(cord_concat, 64, encrypted_len*8);

    i = (t->length/8)-1;
    while(i >= 0) {
        if(t->k[i]) {
        	flag = 1;
        	break;
        }
        --i;
    }
    if(!flag) {
        return -1;
    }

    memcpy(decrypted, encrypted+96, decrypted_len);
    for(i = 0; i < decrypted_len; ++i) {
        decrypted[i] ^= t->k[i];
    }
    #ifdef DEBUG
    printf("C2鈯晅: ");
    print_n_bytes(decrypted, decrypted_len);
    #endif

    memcpy(u, cord_concat, 32);
    memcpy(u+32, decrypted, decrypted_len);
    memcpy(u+32+decrypted_len, cord_concat+32, 32);

    sm3_hash(digest, u, 64+decrypted_len);
    #ifdef DEBUG
    printf("u: ");
    print_n_bytes(digest, 32);
    #endif
    free(u);
    if(memcmp(digest, encrypted+64, 32)!=0) {
        return -1;
    }

    *outputLength = decrypted_len;
    return 0;
}

